use std::{
    cell::RefCell, collections::HashSet, fmt::Debug, hash::Hash, iter::Sum, rc::Rc
};
 
pub use std::ops::{Add, Div, Mul, Neg, Sub};
use super::{op::{BinaryOperate, UnaryOperate}, Operate};

struct VarInner {
    data: f64,
    grad: f64,
    operate: Option<Operate>
}

impl VarInner {
    fn new(data: f64, operate: Option<Operate>) -> Self {
        Self {
            data,
            grad: 0.0,
            operate
        }
    }
}

impl PartialEq for VarInner {
    fn eq(&self, other: &Self) -> bool {
        self.data == other.data &&
        self.grad == other.grad &&
        self.operate == other.operate
    }
}

impl Eq for VarInner {}

impl Debug for VarInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VarInner")
            .field("data", &self.data)
            .field("grad", &self.grad)
            .field("operate", &self.operate)
            .finish()
    }
}


#[derive(Clone, Eq, PartialEq, Debug)]
pub struct Var{
    inner: Rc<RefCell<VarInner>>
}

impl Var {
    pub fn from<T: Into<Self>>(t: T) -> Self {
        t.into()
    }

    fn new(data: f64, operate: Option<Operate>) -> Self {
        Self{ inner: Rc::new(RefCell::new(VarInner::new(data, operate))) }
    }

    pub fn backward(&self) {
        let mut visited = HashSet::new();
        
        fn do_backward(visited: &mut HashSet<Var>, var: &Var) {
            if visited.contains(&var) {
                return;
            }

            visited.insert(var.clone());
            
            if let Some(operate) = &var.inner.borrow().operate {
                match operate {
                    Operate::Binary(bin_op) => {
                        bin_op.backward(var);
                        let (l, r) = bin_op.operands();
                        do_backward(visited, l);
                        do_backward(visited, r);
                    }
                    Operate::Unary(un_op) => {
                        un_op.backward(var);
                        do_backward(visited, un_op.operand());
                    }
                }
            }
        }

        self.inner.borrow_mut().grad = 1.0;
        do_backward(&mut visited, self);
    }

    pub fn pow(&self, power: f64) -> Self {
        let result = self.inner.borrow().data.powf(power);
        Self::new(result, Some(Operate::Unary(UnaryOperate::Pow(self.clone(), power))))
    }

    pub fn tanh(&self) -> Self {
        let result = self.inner.borrow().data.tanh();
        Self::new(result, Some(Operate::Unary(UnaryOperate::Tanh(self.clone()))))
    }

    pub fn exp(&self) -> Self {
        let result = self.data().exp();
        Self::new(result, Some(Operate::Unary(UnaryOperate::Exp(self.clone()))))
    }

    pub fn log(&self, base: f64) -> Self {
        let result = self.data().log(base);
        Self::new(result, Some(Operate::Unary(UnaryOperate::Log(self.clone(), base))))
    }

    pub fn ln(&self) -> Self {
        self.log(1f64.exp())
    }

    pub fn abs(&self) -> Self {
        let result = self.data().abs();
        Self::new(result, Some(Operate::Unary(UnaryOperate::Abs(self.clone()))))
    }

    pub fn sqrt(&self) -> Self {
        self.pow(0.5)
    }

    pub fn relu(&self) -> Self {
        let result = if self.data() < 0. { 0. } else { self.data() }; 
        Self::new(result, Some(Operate::Unary(UnaryOperate::Relu(self.clone()))))
    }

    pub fn sigmoid(&self) -> Self {
        Var::from(1.) / (Var::from(1.) + (-self).exp())
    }

    pub fn data(&self) -> f64 {
        self.inner.borrow().data
    }

    pub fn gard(&self) -> f64 {
        self.inner.borrow().grad
    }

    pub fn increase_data(&self, data: f64) {
        self.inner.borrow_mut().data += data;
    }

    pub fn zero_gard(&self) {
        self.inner.borrow_mut().grad = 0.0;
    }

    pub fn increase_grad(&self, grad: f64) {
        self.inner.borrow_mut().grad += grad;
    }

    pub fn set_grad(&self, grad: f64) {
        self.inner.borrow_mut().grad = grad;
    }

    pub fn adjust(&self, factor: f64) {
        let mut var = self.inner.borrow_mut();
        var.data += factor * var.grad;
    }
}

impl Hash for Var {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        // self.inner.borrow().hash(state);
        self.inner.as_ptr().hash(state);
    }
}

impl<T: Into<f64>> From<T> for Var {
    fn from(value: T) -> Self {
        Self::new(value.into(), None)
    }
}

impl Add<&Var> for &Var {
    type Output = Var;
    fn add(self, rhs: &Var) -> Self::Output {
        let result = self.data() + rhs.data();
        Var::new(result, Some(Operate::Binary(BinaryOperate::Add(self.clone(), rhs.clone()))))
    }
}

impl Add<Var> for Var {
    type Output = Var;
    fn add(self, rhs: Var) -> Self::Output {
        &self + &rhs
    }
}

impl Add<&Var> for Var {
    type Output = Var;
    fn add(self, rhs: &Var) -> Self::Output {
        &self + rhs
    }
}

impl Add<Var> for &Var {
    type Output = Var;
    fn add(self, rhs: Var) -> Self::Output {
        self + &rhs
    }
}

impl Add<f64> for Var {
    type Output = Var;
    fn add(self, rhs: f64) -> Self::Output {
        &self + &Var::from(rhs)
    }
}

impl Add<f64> for &Var {
    type Output = Var;
    fn add(self, rhs: f64) -> Self::Output {
        self + &Var::from(rhs)
    }
}

impl Add<Var> for f64 {
    type Output = Var;
    fn add(self, rhs: Var) -> Self::Output {
        &Var::from(self) + &rhs
    }
}

impl Add<&Var> for f64 {
    type Output = Var;
    fn add(self, rhs: &Var) -> Self::Output {
        &Var::from(self) + rhs
    }
}

impl Mul<&Var> for &Var {
    type Output = Var;
    fn mul(self, rhs: &Var) -> Self::Output {
        let result = self.data() * rhs.data();
        Var::new(result, Some(Operate::Binary(BinaryOperate::Mul(self.clone(), rhs.clone()))))
    }
}

impl Mul<Var> for Var {
    type Output = Var;
    fn mul(self, rhs: Var) -> Self::Output {
        &self * &rhs
    }
}

impl Mul<&Var> for Var {
    type Output = Var;
    fn mul(self, rhs: &Var) -> Self::Output {
        &self * rhs
    }
}

impl Mul<Var> for &Var {
    type Output = Var;
    fn mul(self, rhs: Var) -> Self::Output {
        self * &rhs
    }
}

impl Mul<f64> for Var {
    type Output = Var;
    fn mul(self, rhs: f64) -> Self::Output {
        &self * &Var::from(rhs)
    }
}

impl Mul<f64> for &Var {
    type Output = Var;
    fn mul(self, rhs: f64) -> Self::Output {
        self * &Var::from(rhs)
    }
}

impl Mul<Var> for f64 {
    type Output = Var;
    fn mul(self, rhs: Var) -> Self::Output {
        &Var::from(self) * &rhs
    }
}

impl Mul<&Var> for f64 {
    type Output = Var;
    fn mul(self, rhs: &Var) -> Self::Output {
        &Var::from(self) * rhs
    }
}

impl Sub<&Var> for &Var {
    type Output = Var;
    fn sub(self, rhs: &Var) -> Self::Output {
        let result = self.data() - rhs.data();
        Var::new(result, Some(Operate::Binary(BinaryOperate::Sub(self.clone(), rhs.clone()))))
    }
}

impl Sub<Var> for Var {
    type Output = Var;
    fn sub(self, other: Var) -> Self::Output {
        &self - &other
    }
}

impl Sub<Var> for &Var {
    type Output = Var;
    fn sub(self, other: Var) -> Self::Output {
        self - &other
    }
}

impl Sub<&Var> for Var {
    type Output = Var;
    fn sub(self, other: &Var) -> Self::Output {
        &self - other
    }
}

impl Sub<f64> for Var {
    type Output = Var;
    fn sub(self, rhs: f64) -> Self::Output {
        &self - &Var::from(rhs)
    }
}

impl Sub<f64> for &Var {
    type Output = Var;
    fn sub(self, rhs: f64) -> Self::Output {
        self - &Var::from(rhs)
    }
}

impl Sub<Var> for f64 {
    type Output = Var;
    fn sub(self, rhs: Var) -> Self::Output {
        &Var::from(self) - &rhs
    }
}

impl Sub<&Var> for f64 {
    type Output = Var;
    fn sub(self, rhs: &Var) -> Self::Output {
        &Var::from(self) - rhs
    }
}

impl Div<&Var> for &Var {
    type Output = Var;
    fn div(self, rhs: &Var) -> Self::Output {
        let result = self.data() / rhs.data();
        Var::new(result, Some(Operate::Binary(BinaryOperate::Div(self.clone(), rhs.clone()))))
    }    
}

impl Div<Var> for Var {
    type Output = Var;
    fn div(self, rhs: Var) -> Self::Output {
        &self / &rhs
    }
}

impl Div<Var> for &Var {
    type Output = Var;
    fn div(self, rhs: Var) -> Self::Output {
        self / &rhs
    }
}

impl Div<&Var> for Var {
    type Output = Var;
    fn div(self, rhs: &Var) -> Self::Output {
        &self / rhs
    }
}

impl Div<f64> for Var {
    type Output = Var;
    fn div(self, rhs: f64) -> Self::Output {
        &self / &Var::from(rhs)
    }
}

impl Div<f64> for &Var {
    type Output = Var;
    fn div(self, rhs: f64) -> Self::Output {
        self / &Var::from(rhs)
    }
}

impl Div<Var> for f64 {
    type Output = Var;
    fn div(self, rhs: Var) -> Self::Output {
        &Var::from(self) / &rhs
    }
}

impl Div<&Var> for f64 {
    type Output = Var;
    fn div(self, rhs: &Var) -> Self::Output {
        &Var::from(self) / rhs
    }
}

impl Neg for &Var {
    type Output = Var;
    fn neg(self) -> Self::Output {
        &Var::from(-1) * self
    }
}

impl Neg for Var {
    type Output = Var;
    fn neg(self) -> Self::Output {
        &Var::from(-1) * &self
    }
}

impl Sum for Var {
    fn sum<I: Iterator<Item = Self>>(mut iter: I) -> Self {
        let mut sum = Var::from(0.0);
        loop {
            let val = iter.next();
            if val.is_none() {
                break;
            }

            sum = sum + val.unwrap();
        }
        sum
    }
}

#[cfg(test)]
mod test {
    #[allow(unused)]
    use super::*;

    #[test]
    fn test_new() {
        let v1 = Var::from(1.2);
        let v2 = Var::from(1);
        assert_eq!(v1.data(), 1.2);
        assert_eq!(v2.data(), 1.0);
    }

    #[test]
    fn test_op() {
        let v1 = Var::from(1.2);
        let v2 = Var::from(2.4);
        let v3 = &v1 + &v2;
        let v4 = &v1.add(&v2);
        assert_eq!(v3.data(), 1.2 + 2.4);
        assert_eq!(v4.data(), 1.2 + 2.4);

        let v = Var::from(2.);
        let v = v.pow(2.);
        assert_eq!(4., v.data());

        let v = Var::from(2.).exp();
        assert_eq!(v.data(), 2.0f64.exp());

        let v = Var::from(2.).sigmoid();
        assert_eq!(v.data(), 1. / (1. + (-2.0f64).exp()));
    }

    #[test]
    fn test_backward() {
        // Mul
        let a = Var::from(1.2);
        let b = Var::from(3.7);
        let y = &a * &b;
        y.backward();
        assert_eq!(y.data(), 1.2 * 3.7);
        assert_eq!(a.gard(), 3.7);
        assert_eq!(b.gard(), 1.2);

        let a = Var::from(3.0);
        let y = &a * &a;
        y.backward();
        assert_eq!(y.data(), 3.0 * 3.0);
        assert_eq!(a.gard(), 2.0 * 3.0);
        
        // Add
        let a = Var::from(1.3);
        let b = Var::from(3.4);
        let y = &a + &b;
        y.backward();
        assert_eq!(y.data(), 1.3 + 3.4);
        assert_eq!(a.gard(), 1.0);
        assert_eq!(b.gard(), 1.0);

        // Sub
        let a = Var::from(10);
        let b = Var::from(4);
        let y = &a - &b;
        y.backward();
        assert_eq!(y.data(), 10. - 4.);
        assert_eq!(a.gard(), 1.0);
        assert_eq!(b.gard(), -1.0);

        // Mul
        let a = Var::from(2);
        let b = Var::from(5);
        let y = &a / &b;
        y.backward();
        assert_eq!(y.data(), 2. / 5.);
        assert_eq!(a.gard(), 0.2);
        assert_eq!(b.gard(), -0.08);

        // Comp
        let x1 = Var::from(1.2);
        let x2 = Var::from(2.5);
        let y = 3. * &x1 - 7.1 * &x2;
        y.backward();
        assert_eq!(x1.gard(), 3.);
        assert_eq!(x2.gard(), -7.1);

        let x1 = Var::from(1.2);
        let x2 = Var::from(2.5);
        let y = [&x1 * Var::from(3.), Var::from(-7.1) * &x2].into_iter().sum::<Var>();
        y.backward();
        assert_eq!(x1.gard(), 3.);
        assert_eq!(x2.gard(), -7.1);
    }
}